use std::{collections::HashMap, sync::Arc};
use anyhow::Context;
use tokio::net::TcpListener;
use mcp_sdk::{
    server::{McpServer, McpServerHandler}, 
    transport::tcp::TcpTransport,
    McpClientInitializeResult, 
    McpImplementation, 
    McpPromptsCapability, 
    McpResourcesCapability, 
    McpResult, 
    McpServerCapabilities, 
    McpServerInitializeResult, 
    McpTool, 
    McpToolCallResult, 
    McpToolCallResultContent, 
    McpToolInputSchema, 
    McpToolsCapability
};
use tokio;
use async_trait::async_trait;
use log::{debug, info, LevelFilter};
use simple_logger::SimpleLogger;

struct HelloHandler;

#[async_trait]
impl McpServerHandler for HelloHandler {
    async fn initialize(&self, result: McpClientInitializeResult) -> McpResult<McpServerInitializeResult> 
    {
        info!("Client connected: {} v{}", result.client_info.name, result.protocol_version);
        
        Ok(McpServerInitializeResult {
            protocol_version: mcp_sdk::protocol::LATEST_PROTOCOL_VERSION.to_string(),
            capabilities: McpServerCapabilities {
                experimental: Some(HashMap::new()),
                logging: None,
                prompts: Some(McpPromptsCapability {
                    list_changed: Some(false),
                }),
                resources: Some(McpResourcesCapability {
                    subscribe: Some(false),
                    list_changed: Some(false),
                }),
                tools: Some(McpToolsCapability {
                    list_changed: Some(false)
                })  
            },
            server_info: McpImplementation {
                name: "Hello".to_string(),
                version: "1.0.1".to_string()
            },
            instructions: None
        })
    }

    async fn tools_list(&self) -> McpResult<Vec<McpTool>> {
        Ok(vec![
            McpTool {
                name: "get_my_name".into(),
                description: Some("Return my name to you!".into()),
                input_schema: McpToolInputSchema {
                    r#type: "object".to_string(),
                    required: None,
                    properties: None,
                }
            },
            McpTool {
                name: "say".into(),
                description: Some("print something to stdout!".into()),
                input_schema: McpToolInputSchema {
                    r#type: "object".to_string(),
                    required: Some(serde_json::json!( "value" )),
                    properties: Some(
                        serde_json::json!({
                            "value": {
                                "title": "Value",
                                "description": "the string need to print!",
                                "type": "string"
                            }
                        })
                    ),
                }
            },
        ])
    }

    async fn execute_tool(
        &self,
        tool_name: &str,
        arguments: Option<serde_json::Value>,
    ) -> McpResult<McpToolCallResult> {
        debug!("Received method call: {} with arguments: {:?}", tool_name, arguments);
        if tool_name == "get_my_name" {
            Ok(McpToolCallResult{
                content: vec![ McpToolCallResultContent::Text { text: "molesir".to_string() } ],
                is_error: Some(false)
            })
        } else {
            todo!()
        }
    }

    async fn shutdown(&self) -> McpResult<()> {
        debug!("Handler shutting down");
        Ok(())
    }
}

async fn result_main() -> anyhow::Result<()> {
    SimpleLogger::new().with_level(LevelFilter::Debug).init().context("Init logger")?;

    let listener = TcpListener::bind("0.0.0.0:020716").await?;
    let (stream, addr) = listener.accept().await?;
    info!("New client from {:?}", addr);

    let (transport, _sender) = TcpTransport::new(stream).await;
    let server = McpServer::new(Arc::new(transport), Arc::new(HelloHandler));
    let handle = tokio::spawn(async move {
        info!("Starting server");
        if let Err(e) = server.run().await {
            info!("Error in connection: {}", e);
        }
        info!("Close server");
    });
    handle.await?;

    Ok(())
}

#[tokio::main]
async fn main() {
    if let Err(err) = result_main().await {
        eprintln!("{}", err);
    }
}

